import cv2
import numpy as np
import csv
from datetime import datetime  # 用于添加测量时间戳

# -------------------------- 1. 基础配置（需根据你的图像修改）--------------------------
IMAGE_PATH = "butterfly_with_ruler.jpeg"  # 带参照物的蝴蝶图像路径
REFERENCE_REAL_SIZE = 25.0  # 参照物实际尺寸（mm，如1元硬币直径25mm）
BLUR_KERNEL = (5, 5)  # 高斯模糊核（去噪）
THRESHOLD_VALUE = 127  # 二值化阈值（根据图像亮度调整，100-150常用）
EXPORT_TXT_PATH = "butterfly_measurement.txt"  # TXT结果导出路径
EXPORT_CSV_PATH = "butterfly_measurement.csv"  # CSV结果导出路径

# -------------------------- 2. 核心功能函数 --------------------------
def preprocess_image(img_path):
    """图像预处理：灰度化→模糊→二值化"""
    img_gray = cv2.imread(img_path, 0)
    img_blur = cv2.GaussianBlur(img_gray, BLUR_KERNEL, 0)
    _, img_bin = cv2.threshold(img_blur, THRESHOLD_VALUE, 255, cv2.THRESH_BINARY_INV)
    return img_gray, img_bin

def get_reference_ratio(contours, ref_real_size):
    """识别参照物，计算像素→毫米的换算比例"""
    max_area = 0
    reference_contour = None
    # 找面积最大的轮廓作为参照物（可根据实际调整筛选逻辑）
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area > max_area:
            max_area = area
            reference_contour = cnt
    # 计算参照物外接矩形（取宽高最大值为像素尺寸）
    x_ref, y_ref, w_ref, h_ref = cv2.boundingRect(reference_contour)
    ref_pixel_size = max(w_ref, h_ref)
    pixel_to_mm = ref_real_size / ref_pixel_size  # 换算比例（mm/像素）
    return pixel_to_mm, (x_ref, y_ref, w_ref, h_ref)

def measure_butterfly(contours, reference_contour, pixel_to_mm):
    """测量蝴蝶尺寸：宽度、高度、面积"""
    # 排除参照物轮廓，取剩余最大轮廓为蝴蝶
    butterfly_cnt = max([cnt for cnt in contours if cnt is not reference_contour], 
                        key=cv2.contourArea)
    # 计算外接矩形（宽度=翅膀展开宽，高度=整体高）
    x_bt, y_bt, w_bt, h_bt = cv2.boundingRect(butterfly_cnt)
    # 换算实际尺寸
    real_width = w_bt * pixel_to_mm  # 翅膀展开宽度（mm）
    real_height = h_bt * pixel_to_mm  # 整体高度（mm）
    real_area = cv2.contourArea(butterfly_cnt) * (pixel_to_mm ** 2)  # 总面积（mm²）
    return real_width, real_height, real_area, (x_bt, y_bt, w_bt, h_bt), butterfly_cnt

def export_measurement_data(data, txt_path, csv_path):
    """导出测量数据到TXT和CSV"""
    # 1. 导出TXT（带时间戳和详细说明）
    with open(txt_path, "w", encoding="utf-8") as f:
        f.write(f"蝴蝶翅膀测量报告\n")
        f.write(f"测量时间：{data['timestamp']}\n")
        f.write(f"图像路径：{data['image_path']}\n")
        # 修正：使用正确的键名 reference_real_size_mm
        f.write(f"参照物实际尺寸：{data['reference_real_size_mm']} mm\n")
        f.write(f"像素换算比例：1像素 = {data['pixel_to_mm_ratio']:.3f} mm\n")
        f.write(f"\n===== 测量结果 =====\n")
        f.write(f"翅膀展开宽度：{data['wing_spread_width_mm']:.2f} mm\n")
        f.write(f"蝴蝶整体高度：{data['total_height_mm']:.2f} mm\n")
        f.write(f"蝴蝶总面积（含身体）：{data['total_area_mm2']:.2f} mm²\n")

    # 2. 导出CSV（便于数据统计，字段：时间、图像路径、各测量值）
    # 检查CSV是否已存在，不存在则创建表头
    try:
        with open(csv_path, "r", encoding="utf-8") as f:
            has_header = True
    except FileNotFoundError:
        has_header = False

    with open(csv_path, "a", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=data.keys())
        if not has_header:
            writer.writeheader()  # 首次写入时添加表头
        writer.writerow(data)  # 写入测量数据

# -------------------------- 3. 主流程执行 --------------------------
if __name__ == "__main__":
    # 步骤1：图像预处理
    img_gray, img_bin = preprocess_image(IMAGE_PATH)
    img_color = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2BGR)  # 彩色图用于标注

    # 步骤2：轮廓检测
    contours, _ = cv2.findContours(img_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 步骤3：计算像素-毫米比例
    pixel_to_mm, (x_ref, y_ref, w_ref, h_ref) = get_reference_ratio(contours, REFERENCE_REAL_SIZE)
    # 在图像上标注参照物
    cv2.rectangle(img_color, (x_ref, y_ref), (x_ref + w_ref, y_ref + h_ref), (0, 255, 0), 2)
    cv2.putText(img_color, f"参照物: {REFERENCE_REAL_SIZE}mm", 
                (x_ref, y_ref - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

    # 步骤4：测量蝴蝶尺寸
    real_width, real_height, real_area, (x_bt, y_bt, w_bt, h_bt), butterfly_cnt = \
        measure_butterfly(contours, contours[np.argmax([cv2.contourArea(c) for c in contours])], pixel_to_mm)
    # 在图像上标注蝴蝶和测量结果
    cv2.rectangle(img_color, (x_bt, y_bt), (x_bt + w_bt, y_bt + h_bt), (0, 0, 255), 2)
    cv2.putText(img_color, f"宽度: {real_width:.1f}mm", (x_bt, y_bt - 10), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
    cv2.putText(img_color, f"高度: {real_height:.1f}mm", (x_bt, y_bt - 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
    cv2.putText(img_color, f"面积: {real_area:.1f}mm²", (x_bt, y_bt - 50), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)

    # 步骤5：整理数据并导出
    measurement_data = {
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),  # 测量时间
        "image_path": IMAGE_PATH,  # 图像路径
        "reference_real_size_mm": round(REFERENCE_REAL_SIZE, 2),  # 参照物实际尺寸
        "pixel_to_mm_ratio": round(pixel_to_mm, 3),  # 像素换算比例
        "wing_spread_width_mm": round(real_width, 2),  # 翅膀展开宽度
        "total_height_mm": round(real_height, 2),  # 整体高度
        "total_area_mm2": round(real_area, 2)  # 总面积
    }
    export_measurement_data(measurement_data, EXPORT_TXT_PATH, EXPORT_CSV_PATH)

    # 步骤6：显示结果
    cv2.imshow("Butterfly Measurement (Annotated)", img_color)
    print("="*50)
    print("测量完成！结果如下：")
    print(f"翅膀展开宽度：{real_width:.2f} mm")
    print(f"蝴蝶整体高度：{real_height:.2f} mm")
    print(f"蝴蝶总面积：{real_area:.2f} mm²")
    print(f"\n数据已导出至：")
    print(f"- TXT文件：{EXPORT_TXT_PATH}")
    print(f"- CSV文件：{EXPORT_CSV_PATH}")
    print("="*50)
    cv2.waitKey(0)  # 按任意键关闭图像窗口
    cv2.destroyAllWindows()